package com.ics.atable.chat.memory;

import lombok.Getter;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
import org.springframework.ai.chat.client.advisor.api.BaseChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.util.Assert;
import reactor.core.scheduler.Scheduler;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 自定义Advisor，将对话历史放到系统提示词中
 *
 * @author haohongbin
 */
public final class HistoryAsSystemAdvisor implements BaseChatMemoryAdvisor {
    private final ChatMemory chatMemory;
    private final String defaultConversationId;
    private final int order;
    @Getter
    private final Scheduler scheduler;

    private HistoryAsSystemAdvisor(ChatMemory chatMemory, String defaultConversationId, int order, Scheduler scheduler) {
        Assert.notNull(chatMemory, "chatMemory cannot be null");
        Assert.hasText(defaultConversationId, "defaultConversationId cannot be null or empty");
        Assert.notNull(scheduler, "scheduler cannot be null");
        this.chatMemory = chatMemory;
        this.defaultConversationId = defaultConversationId;
        this.order = order;
        this.scheduler = scheduler;
    }

    @Override
    public int getOrder() {
        return this.order;
    }

    @NotNull
    @Override
    public ChatClientRequest before(ChatClientRequest request, @NotNull AdvisorChain chain) {
        String conversationId = this.getConversationId(request.context(), defaultConversationId);

        // 1) 读取历史，序列化为数组文本（示例使用 JSON 风格字符串）
        List<Message> history = chatMemory.get(conversationId);
        String historyArray = history.stream().map(this::toCompactJson).collect(Collectors.joining(","));
        String historyAsSystemText = "[" + historyArray + "]";

        // 2) 构造“历史数组”的 SystemMessage
        SystemMessage historySystem = new SystemMessage("CONVERSATION_HISTORY=" + historyAsSystemText);

        // 3) 重组本次请求的 messages：
        //    - 保留已有的 system 指令（如果有），再追加 historySystem
        //    - 保留 user 指令等
        List<Message> original = request.prompt().getInstructions();
        List<Message> rebuilt = new ArrayList<>();

        // 3.1 先放已有的 system 指令
        for (Message m : original) {
            if (m instanceof SystemMessage) {
                rebuilt.add(m);
            }
        }
        // 3.2 加入历史数组 SystemMessage
        rebuilt.add(historySystem);

        // 3.3 加入非 system 的其它指令（如 user/tool 等）
        for (Message m : original) {
            if (!(m instanceof SystemMessage)) {
                rebuilt.add(m);
            }
        }

        ChatClientRequest newRequest = request.mutate()
                .prompt(request.prompt().mutate().messages(rebuilt).build())
                .build();

        // 4) 将当前 user 消息写入记忆
        UserMessage userMessage = newRequest.prompt().getUserMessage();
        if (userMessage != null) {
            chatMemory.add(conversationId, userMessage);
        }
        return newRequest;
    }

    @NotNull
    @Override
    public ChatClientResponse after(ChatClientResponse response, @NotNull AdvisorChain chain) {
        String conversationId = getConversationId(response.context(), defaultConversationId);
        List<Message> assistantMessages = new ArrayList<>();
        if (response.chatResponse() != null) {
            assistantMessages = response.chatResponse().getResults().stream()
                    .map(Generation::getOutput).collect(Collectors.toList());
        }
        if (!assistantMessages.isEmpty()) {
            chatMemory.add(conversationId, assistantMessages);
        }
        return response;
    }

    private String toCompactJson(Message m) {
        String role = m.getMessageType().getValue();
        String content = m.getText();
        // 如果是 assistant 消息且包含 output 字段，则提取 output 的值
        if ("assistant".equals(role) && content != null && content.contains("\"output\"")) {
            try {
                // 使用正则表达式提取 output 字段的值
                java.util.regex.Pattern pattern = java.util.regex.Pattern.compile("\"output\"\\s*:\\s*\"([^\"]*)\"");
                java.util.regex.Matcher matcher = pattern.matcher(content);
                if (matcher.find()) {
                    content = matcher.group(1);
                    // 处理转义字符
                    content = StringUtils.replace(content, "\\\"", "\"");
                }
            } catch (Exception e) {
                // 如果解析失败，使用原始内容
                content = m.getText();
            }
        }

        if (content != null) {
            content = StringUtils.replace(content, "\"", "\\\"");
        }
        return role + ":\"" + content + "\"\n";
    }



    public static HistoryAsSystemAdvisor.Builder builder(ChatMemory chatMemory) {
        return new HistoryAsSystemAdvisor.Builder(chatMemory);
    }

    public static final class Builder {
        private String conversationId = "default";
        private int order = -2147482648;
        private Scheduler scheduler;
        private final ChatMemory chatMemory;

        private Builder(ChatMemory chatMemory) {
            this.scheduler = BaseAdvisor.DEFAULT_SCHEDULER;
            this.chatMemory = chatMemory;
        }

        public HistoryAsSystemAdvisor.Builder conversationId(String conversationId) {
            this.conversationId = conversationId;
            return this;
        }

        public HistoryAsSystemAdvisor.Builder order(int order) {
            this.order = order;
            return this;
        }

        public HistoryAsSystemAdvisor.Builder scheduler(Scheduler scheduler) {
            this.scheduler = scheduler;
            return this;
        }

        public HistoryAsSystemAdvisor build() {
            return new HistoryAsSystemAdvisor(this.chatMemory, this.conversationId, this.order, this.scheduler);
        }
    }

}
